import math
import rclpy
from rclpy.node import Node
from rclpy.action import ActionClient, ActionServer
from geometry_msgs.msg import PoseStamped, Quaternion
from std_msgs.msg import Float64MultiArray, Bool
from nav_msgs.msg import Path, Odometry
from rclpy.lifecycle import Publisher
import numpy as np
from rclpy.qos import DurabilityPolicy, QoSProfile
from message_filters import Subscriber, TimeSynchronizer
from general_interface.action import SwitchConnection, AutoDock
import tf_transformations as tft

## Can not work 


class AutoDockTask(Node):

    def __init__(self):
        super().__init__("auto_dock_task_server")

        self.autoDockServer = self.create_action_server(
            AutoDock, "/auto_dock", self.auto_dock_callback
        )

        self.declare_parameter("radius", 0.2)
        self.radius: float = self.get_parameter("radius").value
        self.declare_parameter("step_length", 0.1)
        self.stepLength: float = self.get_parameter("step_length").value
        self.declare_parameter("start_vec", 1.0)
        self.declare_parameter("dock_vec", 0.1)
        self.startVec: float = self.get_parameter("start_vec").value
        self.dockVec: float = self.get_parameter("dock_vec").value
        self.declare_parameter("connection_distance", 1.0)
        self.connectionDistance: float = self.get_parameter("connection_distance").value

    def auto_dock_callback(self, goal_handle):
        self.get_logger().info("auto docking")

        req: AutoDock.Goal = goal_handle.request

        # creat subscription and publisher
        robot_back = req.backward_robot
        robot_forward = req.forward_robot
        qos = QoSProfile(depth=1, durability=DurabilityPolicy.TRANSIENT_LOCAL)
        path_pub1: Publisher = self.create_publisher(
            Path, f"/{self.robot1}/reference_path", qos
        )
        vec_pub1: Publisher = self.create_publisher(
            Float64MultiArray, f"/{self.robot1}/reference_velocity", qos
        )
        path_pub2: Publisher = self.create_publisher(
            Path, f"/{self.robot2}/reference_path", qos
        )
        vec_pub2: Publisher = self.create_publisher(
            Float64MultiArray, f"/{self.robot2}/reference_velocity", qos
        )
        connectorClient = ActionClient(
            self, SwitchConnection, f"/{self.robot1}/switch_connection"
        )

        startDistance = -1.0
        hasStartConnectionTask = False

        def odom_callback(start: Odometry, end: Odometry):
            if startDistance < 0.0:
                startDistance = math.sqrt(
                    (start.pose.pose.position.x - end.pose.pose.position.x) ** 2
                    + (start.pose.pose.position.y - end.pose.pose.position.y) ** 2
                )
            currentDistance = math.sqrt(
                (start.pose.pose.position.x - end.pose.pose.position.x) ** 2
                + (start.pose.pose.position.y - end.pose.pose.position.y) ** 2
            )
            precent =currentDistance / startDistance
            if not hasStartConnectionTask and currentDistance < self.connectionDistance:
                goal = SwitchConnection.Goal()
                goal.command = SwitchConnection.Goal.CONNECT
                connectorClient.send_goal(goal)
                hasStartConnectionTask = True

            path1, vec1 = self.solution1(
                start,
                end,
                self.radius,
                self.startVec,
                self.dockVec,
            )
            path_pub1.publish(path1)
            vec_pub1.publish(vec1)

            path2, vec2 = self.solution1(
                end,
                start,
                self.radius,
                -self.startVec,
                -self.dockVec,
            )
            path_pub2.publish(path2)
            vec_pub2.publish(vec2)

            feedback = AutoDock.Feedback()
            feedback.precent = precent
            feedback.stage = "dock" if hasStartConnectionTask else "approach"
            goal_handle.publish_feedback(feedback)
            if precent >0.95:
                goal_handle.succeed()

        odomSub1 = Subscriber(self, Odometry, f"/{robot_back}/{self.target}")
        odomSub2 = Subscriber(self, Odometry, f"/{robot_forward}/{self.target}")
        odomTs = TimeSynchronizer([odomSub1, odomSub2], 10)
        odomTs.registerCallback(odom_callback)
        
        result = AutoDock.Result()
        result.success = True
        return result

    def solution1(
        self,
        start: Odometry,
        end: Odometry,
        l: float,
        start_vel: float,
        dock_vec: float,
    ) -> tuple[Path, Float64MultiArray]:
        """
        起点和终点以及初始朝向的交点构成三角形，实际轨迹为与三角形相似，但形状更小
        """
        startP = np.array([start.pose.pose.position.x, start.pose.pose.position.y])
        endP = np.array([end.pose.pose.position.x, end.pose.pose.position.y])

        distance = np.linalg.norm(endP - startP)
        l = distance * l

        startQua = (
            start.pose.pose.orientation.x,
            start.pose.pose.orientation.y,
            start.pose.pose.orientation.z,
            start.pose.pose.orientation.w,
        )
        startDir = tft.quaternion_matrix(startQua)[0:2, 0]
        startDir /= np.linalg.norm(startDir)
        endQua = (
            end.pose.pose.orientation.x,
            end.pose.pose.orientation.y,
            end.pose.pose.orientation.z,
            end.pose.pose.orientation.w,
        )
        endDir = tft.quaternion_matrix(endQua)[0:2, 0]
        endDir /= np.linalg.norm(endDir)

        if start_vel < 0:
            startMid = startP - startDir * l
            endMid = endP + endDir * l
            midDir = endMid - startMid
            midDir /= np.linalg.norm(midDir)
            startMidInLine = startMid + midDir * l
            endMidInLine = endMid - midDir * l
            mid = (startMidInLine + endMidInLine) / 2
        else:
            startMid = startP + startDir * l
            endMid = endP - endDir * l
            midDir = endMid - startMid
            midDir /= np.linalg.norm(midDir)
            startMidInLine = startMid + midDir * l
            endMidInLine = endMid - midDir * l
            mid = (startMidInLine + endMidInLine) / 2

        startPoints = self.besizer_curve(
            startP,
            startMid,
            startMidInLine,
            int(np.linalg.norm(startP - startMidInLine) / self.stepLength),
        )
        endPoints = self.besizer_curve(
            endMidInLine,
            endMid,
            endP,
            int(np.linalg.norm(endMidInLine - endP) / self.stepLength),
        )
        midPoints = self.besizer_curve(
            startMidInLine,
            mid,
            endMidInLine,
            int(np.linalg.norm(startMidInLine - endMidInLine) / self.stepLength),
        )
        points = np.vstack([startPoints, midPoints[1:-1, :], endPoints])
        velocity = np.linspace(start_vel, dock_vec, len(points))

        path = Path()
        path.header.frame_id = "map"
        path.header.stamp = self.get_clock().now().to_msg()
        for i in range(len(points)):
            pose = PoseStamped()
            pose.header.frame_id = "map"
            pose.header.stamp = self.get_clock().now().to_msg()
            pose.pose.position.x = points[i][0]
            pose.pose.position.y = points[i][1]
            pose.pose.position.z = 0.0
            if i == len(points) - 1:
                pose.pose.orientation = end.pose.pose.orientation
            else:
                dx = points[i + 1][0] - points[i][0]
                dy = points[i + 1][1] - points[i][1]
                # 插值计算点的朝向
                qua = tft.quaternion_from_euler(
                    0.0,
                    0.0,
                    (
                        math.atan2(
                            -dy,
                            -dx,
                        )
                        if start_vel < 0
                        else math.atan2(dy, dx)
                    ),
                )
                pose.pose.orientation = Quaternion(
                    x=qua[0], y=qua[1], z=qua[2], w=qua[3]
                )
            path.poses.append(pose)
        vec = Float64MultiArray()
        vec.data = velocity
        return path, vec

    def besizer_curve(self, p0, p1, p2, num_points):
        t = np.linspace(0, 1, num_points)[:, np.newaxis]
        points = (1 - t) ** 2 * p0 + 2 * t * (1 - t) * p1 + t**2 * p2
        return points


def main(args=None):
    rclpy.init(args=args)
    task = AutoDockTask()
    rclpy.spin(task)
    rclpy.shutdown()


if __name__ == "__main__":
    main()
